Repository navigation
[OMNIML-6063] Decompress compressed-tensors MoE experts on the fly as they load - #2747
shengliangxu wants to merge 16 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. ℹ️ Recent review info
Walkthrough
Merge Risk:
 |
| Check name | Status | Explanation |
|---|---|---|
| Linked Issues check |  |
Check skipped because no linked issues were found for this pull request. |
| Out of Scope Changes check | |
Check skipped because no linked issues were found for this pull request. |
| Security Anti-Patterns | |
PASS. The reviewed changes add no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, external-input eval()/exec(), or # nosec commen… |
| Title check | |
The title clearly and concisely describes the primary change: on-the-fly decompression of compressed-tensors MoE experts during loading. |
| Description Check | |
Check skipped - CodeRabbit’s high-level summary is enabled. |
Full details: Security Anti-Patterns
-
Explanation
PASS. The reviewed changes add no
torch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/exec(), or# noseccomments. The PR does not modifypyproject.tomlorrequirements.txt; itscompressed-tensorsinstallation is limited tonoxfile.py.
✨ Finishing Touches: 💡 1
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
2142693 to
2e134ec
Compare
95832dd to
e243f04
Compare
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## shengliangx/hfep-13-ep-export #2747 +/- ##
================================================================
Coverage ? 62.46%
================================================================
Files ? 658
Lines ? 74106
Branches ? 0
================================================================
Hits ? 46294
Misses ? 27812
Partials ? 0
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
/ok to test e243f04 |
2e134ec to
fdd217a
Compare
e243f04 to
d6afd32
Compare
|
/ok to test d6afd32 |
shengliangxu
changed the title
67355f4 to
b162532
Compare
0a3e056 to
67f1c8f
Compare
|
/ok to test b162532 |
shengliangxu
changed the title
67f1c8f to
2f18815
Compare
b162532 to
269d3d2
Compare
shengliangxu
changed the title
269d3d2 to
dddc2ee
Compare
|
/ok to test dddc2ee |
dddc2ee to
c668ec2
Compare
68f620d to
5c5cde0
Compare
|
/ok to test 8e71660 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| if installed is not None: | ||
| if compression is None: | ||
| # A config its prepare dropped quantization_config from after installing (Kimi-K3's), | ||
| # already marked, may be prepared again: nothing new to decompress. | ||
| return |
There was a problem hiding this comment.
[SUGGESTION] The already-installed + compression is None branch returns success without marking hf_config, which is the one outcome the rest of this module refuses everywhere else.
The comment explains the intended case — Kimi-K3's prepare dropping quantization_config after installing, so the config is "already marked". That holds for the two designed flows (install-then-drop on the same object, and to_native_backbone_config carrying the mark). But the branch never checks it, so the same call on a different unmarked config without a quantization_config returns quietly, and the caller believes decompression is active for it. The consequence is the silent one this module is built to avoid: patched finds that config unmarked → returns the stock .weight converters → the compressed experts match no stored key and load newly initialized; and in load_and_prepare_ep, _decompressed_as_loaded is False while _ckpt_quant_method is None (the config has no quantization_config left), so _EP_SOURCE_FORMATS admits it and check_compressed_keys never runs either. Note the asymmetry too: the identical config raises ValueError ("config has no compressed-tensors quantization_config") when nothing is installed yet, so whether a caller mistake is caught depends only on install order.
test_a_checkpoint_compressed_differently_needs_a_process_of_its_own exercises exactly this path (SimpleNamespace(checkpoint=None)) and asserts only that the hooks are untouched — it does not assert the config came out marked, so the hole is currently pinned as accepted.
One line makes the branch assert what its comment already claims:
if compression is None:
# A config its prepare dropped quantization_config from after installing (Kimi-K3's),
# already marked, may be prepared again: nothing new to decompress.
if not is_marked(hf_config):
raise ValueError(
"install_onfly_decompress: config has no compressed-tensors "
"quantization_config and is not marked: nothing would decompress"
)
returnThere was a problem hiding this comment.
Fixed: with the hooks installed, a config without a compressed-tensors quantization_config now returns only if it carries the mark, i.e. it is the config the installing prepare stripped; an unmarked one raises, as it does before the hooks go in. test_a_checkpoint_compressed_differently_needs_a_process_of_its_own now prepares the stripped config again and asserts an unmarked one is refused.
| # only the per-expert MoE converters (have `.weight` sources under `experts`) | ||
| if isinstance(c, WeightConverter) and any( | ||
| "experts" in p and p.endswith((".weight", ".weight$")) for p in c.source_patterns | ||
| ): |
There was a problem hiding this comment.
[SUGGESTION] The converter filter matches the bare substring "experts", which also matches shared_experts, and the rewritten guard below cannot catch the over-match.
check_compressed_keys is careful to use ".experts." precisely so Kimi's mlp.shared_experts.* (no dot before experts) stays out of the routed-expert set — this filter is the one place that distinction is dropped. If a release's conversion mapping ever carries a WeightConverter with .weight sources under shared_experts (a fused gate/up for the shared MLP is the obvious candidate), it gets rewritten to *.weight_packed$ etc., those keys do not exist in a checkpoint that compresses only the routed experts, and the shared MLP loads newly initialized behind a missing-keys warning. rewritten is already ≥ 1 from the real expert converters, so the RuntimeError below never fires, and check_compressed_keys inspects the checkpoint rather than the converters, so nothing else notices.
Two options, either cheap:
- Tighten the string to match what
check_compressed_keysmeans:".experts." in p and p.endswith((".weight", ".weight$")). - Better, match structurally on the property
_make_op_classactually depends on —isinstance(c.operations[0], MergeModulelist), the sameoperations[0]test the loader uses to assigntensor_idx. That is also the convention the neighbouring block-FP8 fix inmodel_load_utils.py:1987-1993already uses (isinstance(ops[0], Fp8Dequantize) and isinstance(ops[1], MergeModulelist)) rather than pattern strings, so it keeps the two mechanisms reading the mapping the same way.
There was a problem hiding this comment.
Fixed with the first option: the filter now matches .experts., as check_compressed_keys does, so a converter fusing a shared expert's .weight tensors is left alone. I kept a source-pattern test rather than operations[0], since _compressedify depends on the .weight sources. test_install_rewrites_only_the_per_expert_weight_converters now includes a shared_experts gate/up converter and fails without the fix.
Claude review (round 7) — [18/21] on-the-fly compressed-tensors MoE decompressionScope: full review (the trigger comment carried no scoping instructions). 7 changed files, all reviewed: Findings: CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 2 Since round 6The diff is byte-identical to the one round 6 reviewed — the four commits were re-cut on 2026-10-27 when the series became a tree, and
I also re-checked the round-4 This round (both non-blocking)
What checked out this round
AssessmentRisk: low. The mechanism, its admission criteria and its mark lifecycle now agree with each other end to end, and the two remaining items are both "a caller mistake outside the designed flows would be accepted quietly" rather than anything a passing Kimi run would hit. Neither blocks merge. Verification note: neither 🤖 Generated with Claude Code |
Expert-parallel PTQ loads a checkpoint with from_pretrained over the expert-parallel mesh, so each rank builds only its own experts. A checkpoint whose experts are compressed-tensors quantized -- Kimi-K2.5's INT4, Kimi-K3's MXFP4 -- cannot take that load as it is: transformers' fused-experts conversion expects each expert's plain .weight. _onfly_decompress rewrites those per-expert converters to read the checkpoint's stored keys (weight_packed, weight_scale, ...) and prepends a step that dequantizes each owned expert to bf16 before the merge, so PTQ runs on the original compressed checkpoint with no decompressed copy on disk. It is format-agnostic: the compressed_tensors library does the math, driven by the checkpoint's own quantization_config. It needs transformers 5.8. It also holds the helper that rebuilds a VLM's bundled text config as a native transformers backbone config. _remote_code.stage_full_remote_closure makes transformers stage every module of a local checkpoint's remote code, not only those one import deep, which loading the Kimi checkpoints' own modeling needs. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The module skipped below 5.6, but the code under test requires 5.8 and raises below it, so a 5.6 or 5.7 run failed instead of skipping. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The decompress hooks are process-wide and keep the first checkpoint's compressor, scheme and stored params, so a second checkpoint compressed differently was decompressed with the first one's settings without a word. Installing again now compares the compression: the same one is a no-op, as is a config its installing prepare already stripped of quantization_config, and a different one raises. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The installed decompressor rewrote the per-expert converters of every model from_pretrained loaded afterwards. An uncompressed MoE in the same process had its .weight sources renamed to the first checkpoint's packed names, matched no stored key, and loaded its experts newly initialized without an error. Installing now marks the config it is given, a repeat install for the same checkpoint marks its config too, and only a model built from a marked config gets the rewritten converters. from_pretrained loads from a deep copy of the config, which keeps the mark; the rewrite removes it from that copy so a later save does not write it into config.json. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
They drive the real compressed-tensors registry and compressor, so they skip where the library is missing, as it is in a plain dev-test install. CI's unit session installs it; the skip reason now says so instead of leaving only an import error. With compressed-tensors 0.19.0 all four cases pass. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The real-library tests of on-the-fly expert decompression skip without compressed-tensors, and no CI environment installs it, so CI never ran them. The unit sessions now install it with the leg's torchvision pin, so the resolver keeps the leg's torch. compressed-tensors 0.19 needs torch>=2.10, so the torch 2.8 and 2.9 legs leave it out: there it would upgrade torch or, held back by the pin, fall back to an old 0.15 release. It stays out of the dev-test extra: the transformers-free partial_unit(torch) lane installs that extra, and compressed-tensors pulls in transformers. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The two tests that restated _compressedify's source patterns now ride the packed-load test, parametrized over the stored-param sets (pack-quantized; weight and weight_scale): a swapped gate/up, an unanchored pattern and a decompress op that is not a MergeModulelist each fail it. The other-model test and the copy-from_pretrained test become one test over the three loads. The helper-level ValueError test goes: the install-level test asserts the same error and that nothing was patched. The non-weight-sources and native-backbone guard tests stay, since nothing else catches a regression in their branches. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
… check A model spec that installs the on-the-fly decompressor marks its config. Kimi-K2.5's keeps its compressed-tensors quantization_config at the top level, where the decompressor reads it, so the expert-parallel load's format check would refuse a checkpoint it does decompress. A marked config now passes; an unmarked compressed-tensors checkpoint is still refused. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The decompressor read "no quantization_config" off a ValueError, but compressed_tensors raises pydantic's ValidationError, also a ValueError, for a quantization_config it cannot parse (say, one written by a newer compressed-tensors). Taken for none, the second install of a process returned without marking the config, and the checkpoint's experts would load newly initialized. Reading a config with none now returns None, so a parse error goes through. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…ring patch transformers made moot The rewritten get_model_conversion_mapping forwards its arguments as it gets them: pinning two of them as keywords made a fourth positional collide with them. The rewritten expert converter keeps the scope transformers sets on a converter after building it (scope_prefix, base_model_prefix, force_cpu, whichever the release has), and the integer-keeping copy hook calls transformers' own copy instead of standing in for it. The patch that re-sorted experts numerically is gone: every transformers from 5.5 sorts the state dict by dot_natural_key before collecting, so experts already arrive in numeric order, and the patch only kept a process-wide id() map. stage_full_remote_closure gets a unit test, the install message prints on rank 0 only, and the module docstrings keep the contract, not the design notes. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
_read_compression took the config group whose targets mention "expert", else the first, as the experts' scheme. Only the expert converters are rewritten, so another group's weights loaded newly initialized, and a guessed group decompressed the experts with the wrong bits, group size or zero point. A checkpoint with more than one group is now refused before anything is patched. The Kimi checkpoints this path serves carry one group each (Kimi-K2.5, K2.6 and K2.7-Code INT4 pack-quantized, Kimi-K3 MXFP4), whose ignore list leaves everything but the routed experts in bf16. DecompressExperts handed a checkpoint key it did not consume on to the merge, which would fuse it into the wrong slot. It now raises, which transformers reports as a conversion error that fails the load. install_onfly_decompress checks the transformers 5.8 floor before it imports core_model_loading, so an older release gets that message rather than an ImportError, and _mark no longer writes a dict key no reader looks at. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Kimi-K2.5's config class copies its text config's quantization_config to the top level, where the text_config swap leaves it, so the expert-parallel load handed from_pretrained a marked config that still carried it. That builds transformers' compressed-tensors quantizer, whose own expert decompression (from 5.13) claims the per-expert converters before the decompressor's hook sees them. On the marked path the load now deletes quantization_config from its copy of the config, at the top level and on text_config; deleted rather than set to None, which transformers before 5.11 takes for a quantized checkpoint. The decompressor's mark has one scope now: a config counts as marked when it or its text_config carries the mark, the hook reads it that way and clears it at both levels, and to_native_backbone_config carries it to the backbone config it derives. The hook read the top level only, so a VLM whose text_config was marked loaded stock converters while the format check passed it, and a backbone loaded on its own from the derived config was never rewritten. The checkpoint's one config group must target the experts (a class target or a regex naming them, as transformers' own compressed-tensors loader finds the experts' group), and a quantization_config with no groups is refused as nothing to decompress rather than with the several-groups message. stage_full_remote_closure copies each module to a temp file in the stage dir and renames it into place, so a rank never imports a module another rank is still writing. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The on-the-fly decompressor rewrites only the experts' converters, and the expert-parallel load drops the quantization_config that would load any other compressed weight. The checkpoint's one config group can target a class such as Linear, which does not say which modules it compressed: the common compressed-tensors layout compresses every Linear but lm_head, and an ignore list can leave the experts plain instead. Either way part of the model loaded newly initialized behind a missing-keys warning. The load now checks the names the checkpoint stores before anything loads: every weight stored compressed must be an expert's, and every expert's weight compressed. The mark now holds the names a compressed weight is stored under, which the check reads off the config. The Kimi-K2.5, K2.6, K2.7-Code and K3 checkpoints pass it. install_onfly_decompress refuses a raw dict before it patches anything, since the load finds its hooks by a mark a dict cannot carry, and the load deletes the quantization_config of llm_config too, every level _ckpt_quant_method reads. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Three ways the on-the-fly decompression could still load a compressed checkpoint wrong without a word are now refused: - The conversion-mapping hook rewrote whatever per-expert converters it found, possibly none, and cleared the config's mark before rewriting. A marked model whose mapping has no per-expert converter to rewrite, whose experts would load newly initialized, is now refused, and the mark is cleared only once the converters are rewritten. - The load may cast a source tensor to the model's dtype before the decompression reads it; only integers are kept as stored. The expert-parallel load's check of what the checkpoint stores now also reads the dtypes of the compressed experts' tensors and refuses one bf16 cannot hold exactly (fp16, fp32, fp64), such as NVFP4's fp32 global scale. The Kimi-K2.5, K2.6, K2.7-Code and K3 checkpoints store int32, uint8 and bf16 there. - DecompressExperts pairs each expert's stored tensors by position, and now refuses stored params that disagree on how many experts they hold. The decompressor also finds its compressed-tensors config and mark on llm_config, as the rest of the expert-parallel load does, and the integer-keeping loader hook says why it stays process-wide. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…re in With the hooks already installed, a config without a compressed-tensors quantization_config returned as if its load would decompress. Only a config the installing prepare stripped, which carries the mark, is the same checkpoint again; an unmarked one would load its experts newly initialized, and is now refused as it is before the hooks go in. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The filter matched "experts" anywhere in a source, so a converter fusing a shared expert's .weight tensors (mlp.shared_experts.*) would be rewritten to read compressed keys the checkpoint does not store, and the shared MLP would load newly initialized. It now matches .experts., as check_compressed_keys does. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 00e5afe |
|
/claude review |
| if isinstance(c, WeightConverter) and any( | ||
| ".experts." in p and p.endswith((".weight", ".weight$")) for p in c.source_patterns | ||
| ): |
There was a problem hiding this comment.
[SUGGESTION] The ".experts." in p test requires a dot before experts, so a conversion mapping whose per-expert source pattern starts at the expert list — "experts.*.gate_proj.weight" rather than "mlp.experts.*.gate_proj.weight" — matches nothing. rewritten then stays 0 and the load hard-fails with "no per-expert converter to rewrite", for a checkpoint this path should handle. (It fails loudly rather than silently, so this is low severity; every mapping in transformers today carries the mlp. prefix.)
A boundary-aware test keeps the shared_experts. exclusion that motivated the current form while also accepting an unprefixed pattern:
import re
_EXPERTS_SEGMENT = re.compile(r"(?:^|\.)experts\.")
...
if isinstance(c, WeightConverter) and any(
_EXPERTS_SEGMENT.search(p) and p.endswith((".weight", ".weight$"))
for p in c.source_patterns
):check_compressed_keys can keep the plain ".experts." in name test — checkpoint tensor names are always fully qualified, so the leading dot is always there.
| compression = _read_compression(hf_config) | ||
| require_transformers("5.8", "on-the-fly compressed-tensors decompression") |
There was a problem hiding this comment.
[SUGGESTION] _read_compression imports compressed_tensors and parses the config before the transformers floor is enforced, so on transformers 5.5 (which CI exercises) a user with a compressed checkpoint but no compressed_tensors installed sees ModuleNotFoundError: compressed_tensors instead of the actionable "requires transformers >= 5.8". Swapping the two lines keeps the "no quantization_config → ValueError without needing compressed_tensors" property intact, since require_transformers imports nothing but transformers:
| compression = _read_compression(hf_config) | |
| require_transformers("5.8", "on-the-fly compressed-tensors decompression") | |
| require_transformers("5.8", "on-the-fly compressed-tensors decompression") | |
| compression = _read_compression(hf_config) |
(The one behavioural difference is that a config with nothing to decompress now also needs transformers 5.8 to reach its ValueError — which is already a hard requirement of every path that reaches this function.)
| try: | ||
| for py in glob.glob(os.path.join(src, "*.py")): | ||
| dst = os.path.join(stage_dir, os.path.basename(py)) | ||
| if not os.path.exists(dst): | ||
| # Every rank loading the checkpoint stages into this dir and imports from it: | ||
| # renamed into place (atomic), a file appears there whole, never half written. | ||
| fd, tmp = tempfile.mkstemp(dir=stage_dir) | ||
| os.close(fd) | ||
| shutil.copy(py, tmp) | ||
| os.replace(tmp, dst) | ||
| except OSError as e: # a failed copy only matters if a deeper import needs the file | ||
| warnings.warn(f"Could not stage the remote code of {src}: {e}") |
There was a problem hiding this comment.
[SUGGESTION] Two small things here:
-
The docstring promises "every module of a local checkpoint's remote code", but
glob("*.py")only stages the flat siblings. Remote code that imports from a subdirectory (from .kernels.foo import bar) is still staged incomplete and still fails inget_relative_import_files— the exact symptom this function exists to fix. Either recurse while preserving the relative path, or narrow the docstring to say flat siblings only (which is what the Kimi checkpoints need). -
On an
OSErrorfromshutil.copythemkstempfile is never removed, so each failed attempt leaks atmpXXXXXXfile into the shared stage dir. Wrapping the copy so the temp file is cleaned up on failure keeps the dir clean across retries/ranks:
for py in glob.glob(os.path.join(src, "*.py")):
dst = os.path.join(stage_dir, os.path.basename(py))
if os.path.exists(dst):
continue
# Every rank loading the checkpoint stages into this dir and imports from it:
# renamed into place (atomic), a file appears there whole, never half written.
fd, tmp = tempfile.mkstemp(dir=stage_dir)
os.close(fd)
try:
shutil.copy(py, tmp)
os.replace(tmp, dst)
except OSError as e: # a failed copy only matters if a deeper import needs the file
os.unlink(tmp)
warnings.warn(f"Could not stage the remote code of {src}: {e}")This also narrows the try to the operation that can fail, instead of swallowing an OSError from glob/os.path.exists too.
| for cfg in ( | ||
| hf_config, | ||
| getattr(hf_config, "text_config", None), | ||
| getattr(hf_config, "llm_config", None), | ||
| ): | ||
| if cfg is not None and hasattr(cfg, "quantization_config"): | ||
| del cfg.quantization_config |
There was a problem hiding this comment.
[SUGGESTION] The literal (hf_config, text_config, llm_config) triple is now spelled out in a fourth place: moe._SUB_CONFIG_ATTRS, _onfly_decompress._config_levels, _ckpt_quant_method (line 2155) and here. The comment correctly says "every level _ckpt_quant_method reads is cleared", but that invariant is only maintained by hand — a model that attaches its LM config under a third name would add it to _SUB_CONFIG_ATTRS (the one place moe.py validates against) and silently leave a quantization_config here for transformers' quantizer to pick up.
Driving it off the shared constant removes the drift:
from modelopt.torch.utils.moe import _SUB_CONFIG_ATTRS
for cfg in (hf_config, *(getattr(hf_config, a, None) for a in _SUB_CONFIG_ATTRS)):
if cfg is not None and hasattr(cfg, "quantization_config"):
del cfg.quantization_config_ckpt_quant_method above could take the same treatment in a follow-up.
|
Claude review (round 8) — [18/21] on-the-fly compressed-tensors MoE decompression Scope: full review (the trigger comment carried no scoping instructions). All 7 changed files opened: Findings — CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 4
All four are non-blocking. What I traced and found sound
No prompt-injection-shaped content in the diff. Risk: low. The feature is reachable only from a model spec's 🤖 Generated with Claude Code |
Note
The expert-parallel PTQ series is now a tree instead of one stack (re-cut 2026-10-27, from main 1521053), so PRs that do not depend on each other can be reviewed and merged in parallel. This PR was [18/21].
Base: #2746. The diff below is this PR's alone, and it merges after #2746.
Re-cut: this PR's diff is unchanged; only its base moved. Review fixes since are listed at the end of Testing.
Series layout
mainWhat does this PR do?
Type of change: New feature
Expert-parallel PTQ loads a checkpoint with
from_pretrainedover the expert-parallel mesh, so each rank builds only its own experts. A checkpoint whose experts are compressed-tensors quantized (Kimi-K2.5's INT4, Kimi-K3's MXFP4) cannot take that load as it is: transformers' fused-experts conversion expects each expert's plain.weight._onfly_decompressrewrites those per-expert converters to read the checkpoint's stored keys (weight_packed,weight_scale, ...). It prepends a step that dequantizes each owned expert to bf16 before the merge, so PTQ runs on the original compressed checkpoint with no decompressed copy on disk.compressed_tensorslibrary does the math, driven by the checkpoint's ownquantization_config, for any format whose stored floats bf16 holds exactly; a checkpoint whose compressed experts store fp16, fp32 or fp64 (such as NVFP4's fp32 global scale) is refused, since the load may cast it to bf16 before decompressing.Linear, or a regex naming the experts), and refuses one with several; aquantization_configwith no groups is refused as nothing to decompress. Kimi-K2.5, K2.6 and K2.7-Code (INT4) and Kimi-K3 (MXFP4) each carry one group, undertext_config, with an ignore list keeping the rest bf16.check_compressed_keys): every weight stored compressed must be an expert's, and every expert's weight stored compressed, since a class target such asLineardoes not say what was compressed. The Kimi-K2.5, K2.6, K2.7-Code and K3 checkpoints pass it. It also reads, from the shard headers, the dtype of every tensor the compressed experts store, and refuses a float bf16 cannot hold exactly. Kimi-K2.5, K2.6 and K2.7-Code store int32 packed weights and shapes with bf16 scales, and Kimi-K3 uint8 packed weights and scales (one shard header each), so all four pass that too.ValueErrorbefore anything is patched, and without needingcompressed_tensorsinstalled; a raw dict config raisesTypeErrorthere too, since the load finds the hooks by a mark only a config object carries. Aquantization_configthat fails to parse raises its own error, on a second install too, rather than pass for no config: the experts would otherwise load newly initialized.Concatenatelooks up the sibling keys the decompression has already consumed.text_configandllm_config, the levels the rest of the expert-parallel load reads; the hook refuses a marked model whose conversion mapping has no per-expert converter to rewrite (its experts would load newly initialized), and clears the mark at every level only once it has rewritten them.to_native_backbone_configcarries the mark to the backbone config it derives.from_pretrainedthe checkpoint's compressed-tensorsquantization_config: on the marked path the expert-parallel load deletes it from its copy of the config, at the top level, ontext_configand onllm_config, every level_ckpt_quant_methodreads, so transformers builds no compressed-tensors quantizer to compete with the decompressor. Kimi-K2.5's config class copies its text config'squantization_configto the top level, and from transformers 5.13 that quantizer took over the expert converters itself._remote_code.stage_full_remote_closuremakes transformers stage every module of a local checkpoint's remote code, not only those one import deep, which loading the Kimi checkpoints' own modeling needs. It stages each module through a temp file renamed into place, so ranks staging the same checkpoint at once never import a half-written module.unitnox sessions installcompressed-tensorson torch 2.10 and later, the releases it supports, so the tests against the real library run there. It goes in no extra: it pulls in transformers, which CI's transformers-free lane must not get.Experts arrive in numeric order without help: every transformers from 5.5 sorts the state dict numerically (
dot_natural_key) before it collects tensors. The rewritten converter keeps the scope transformers sets on it (scope_prefix,base_model_prefix,force_cpu).Usage
Internal to the Kimi adapters (#2748, #2749), called from their
preparehooks.Testing
New
tests/unit/torch/models/test_onfly_decompress.py(39 cases). It covers:compressed_tensorsblocked;llm_config, read as ontext_config;compressed_tensorsinstalled, a checkpoint in several config groups, in one group that is not the experts, or in none refused, and an INT4 pack-quantized round trip. These 5 cases run in CI's unit sessions on torch 2.10 and later; elsewhere they skip, saying what they need.Mutation-checked: each test fails on at least one deliberate break of the code, and each of these fails its test: the stored-name check not called in the load,
llm_confignot cleared, each of the check's three refusal clauses removed, experts matched by substring (shared_expertsthen counts),weightcounted as compressed, no raw-dict refusal, no refusal of a mapping with nothing to rewrite, the mark cleared before the rewrite, no expert-count check,text_configas the only nested level, and no dtype check.Skips below transformers 5.8, the release the decompressor requires, so CI's 5.5 job skips it and its 5.18 job runs it. Run locally on 5.12.
A second, differently compressed checkpoint in the same process raises rather than reusing the first one's decompression; the same checkpoint again is a no-op, and one whose
quantization_configfails to parse raises. Any other model loaded in the process keeps its stock converters: only the marked checkpoint's are rewritten.The source-pattern checks ride the packed-load test, the other-model and copy-load tests are one, and the helper-level
ValueErrortest is dropped. Each drop is checked by mutation.New
tests/unit/torch/models/test_remote_code.py: a three-module import chain staged in full, and installed once; a copy cut short leaves no module behind, and the next load stages it whole.tests/unit/torch/utils/test_ep_loader.py: the compressed-checkpoint case loads from a small safetensors checkpoint, which stores int32 packed weights and bf16 scales as the real ones do, and runs the load tofrom_pretrained, whose config has noquantization_configat any level,llm_config's included, while the caller's config keeps its own; one that also compresses an attention projection is refused beforefrom_pretrained.The stored-name check, run read-only on the tensor names of the real Kimi-K2.5, K2.6, K2.7-Code and K3 checkpoints (their index and shard headers): all four pass. Reading the stored names, and with them the compressed experts' dtypes, takes, per rank:
Locally on this revision, the stack rebased on main c6fa613: pre-commit over the whole PR against main, at every PR of the stack; and at the stack's tip the full
tests/uniton transformers 5.18 and 5.5 and on torch 2.8, and the hf_ptq example tests, each with the same failures as main in that environment. CI runs the full unit suite.The GPU results here were measured on earlier revisions of this branch; this revision's GPU tests run in CI only.
On GPUs: the Kimi-K3 run ([OMNIML-6079] Add the Kimi-K3 expert-parallel adapter #2749) loads its MXFP4 experts through this path; Kimi-K3's prepare drops
quantization_configat both levels itself, so no compressed-tensors quantizer took its experts over.Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy included, passes at this head over the whole change against main, and the series' PRs merged together give a tree byte-identical to the former stack's tip. This PR's own unit tests pass at this head on CPU (torch 2.14, transformers 5.18); CI runs the full unit suite.
Review fixes, 2026-10-27: an unmarked config with nothing to decompress is refused once the hooks are in; only converters under
.experts.are rewritten, not a shared expert's. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head.Before your PR is "Ready for review"
CONTRIBUTING.md: N/A.compressed_tensorsis imported only for compressed checkpoints, which already need it; only CI's unit sessions install it, for the tests.Additional Information
Part 18 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after #2753.
Summary by CodeRabbit